import re

from PyPDF2.generic import Fit
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer, LTContainer, LTChar, LTTextBoxHorizontal
from PyPDF2 import PdfWriter, PdfReader

class PdfUtil:

    @staticmethod
    def log(msg):
        print(msg)

    @staticmethod
    def get_pdf_page_count(pdf_file_path):
        with open(pdf_file_path, 'rb') as file:
            pdf_reader = PdfReader(file)
            page_count = len(pdf_reader.pages)
        return page_count

    @staticmethod
    def extract_text_from_element(element: []):
        """从给定的元素提取文本。"""
        if not isinstance(element, LTTextBoxHorizontal):
            return ""

        text = ""
        for text_line in element:
            # 处理文本行中的每个字符
            for character in text_line:
                if isinstance(character, LTChar):
                    text += character.get_text()
        # text = PdfUtil.get_text_from_text_box(element)
        return text.strip()

    @staticmethod
    def extract_headings(pdf_path, patten, begin_index=0, progress_callback=None):
        headings = []
        edge_x = 1000
        last_title = 'XXX'
        # predicted_next_title_num = None
        pages = []
        page_total = PdfUtil.get_pdf_page_count(pdf_path) - begin_index
        for page_number, page_layout in enumerate(extract_pages(pdf_path), start=1):
            if page_number < begin_index:
                continue
            pages.append(page_layout)
            if progress_callback:
                progress_callback(page_total, page_number - begin_index)

            page_layout = sorted(page_layout, key=lambda x: (-x.y0, x.x0))
            for element in page_layout:
                text = PdfUtil.extract_text_from_element(element).strip()
                # if element.x0 > edge_x + 50:
                #     continue
                if re.match(r'^\d *[)）>]', text):
                    continue
                if re.match('^\d{4,}', text):
                    continue
                if not (text and 4 <= len(text) < 70 and re.match(patten, text)):
                    last_title = text
                    continue
                # elif predicted_next_title_num and text.count('.') > 0:
                #     in_possible = False
                #     for possible_title_num in predicted_next_title_num:
                #         if text.startswith(possible_title_num):
                #             in_possible = True
                #             break
                #     if not in_possible:
                #         last_title = text
                #         continue


                # 去除连续中文之间的空格
                text = re.sub(r'(?<=\w)\s(?=\w)', '', text)
                text = re.match(patten, text).group()

                # 标题可能的格式
                #    3.2.2.2.1 物理信号
                #    3.2.2.2.1物理信号
                #    3.2.2.2.1. 物理信号
                #    3.2.2.2.1.物理信号
                # 提取标题的数字前缀，如 3.2.2.2.1
                title_num = re.match(r'^[1-9]{1,2}(\.\d{1,2})*', text).group()
                heading_level = title_num.count('.') # 书签层级，从 0 开始
                if title_num == text:
                    suffix = ''
                else:
                    suffix = text[len(title_num):].strip()
                    if suffix[0] == '.' or suffix[0] == '、':
                        suffix = suffix[1:]
                text = title_num + ' ' + suffix

                if (len(headings) > 2 and headings[-1]['level'] > 0
                        and PdfUtil.compare_title_num(title_num, headings[-1]['title_num']) < 0):
                    continue
                PdfUtil.log(f'[{page_number} - ({element.x0}, {round(element.y0, 2)}) - {heading_level}] [{text}]')

                if element.x0 < edge_x:
                    edge_x = element.x0

                headings = PdfUtil.remove_prev_mistake_bookmarks(headings, title_num)

                # 增加兼容性，假如识别到的标题为依次为：
                #    1.1 abc
                #    1.2 def
                #    1.3 ghi
                #    2.1 jkl
                #    2.2 mno
                # 则自动增加父级标题 1 xxx, 2 xxx
                if len(headings) == 0:
                    if title_num == '1.1':
                        headings.append(
                            {
                                'text': f'1 {last_title}',
                                'page': page_number,
                                'x': element.x0,
                                'y': element.y0 + 30,
                                'level': 0,
                                'title_num': '1'
                            }
                        )
                elif title_num.endswith('.1') and title_num[:title_num.rfind('.')] != headings[-1]['title_num']:
                    parent_title_num = title_num[:title_num.rfind('.')]
                    headings.append(
                        {
                            'text': f'{parent_title_num} {last_title}',
                            'page': page_number,
                            'x': element.x0,
                            'y': element.y0 + 30,
                            'level': heading_level - 1,
                            'title_num': parent_title_num
                        }
                    )

                headings.append(
                    {
                        'text': text,
                        'page': page_number,
                        'x': element.x0,
                        'y': element.y0,
                        'level': heading_level,
                        'title_num': title_num
                    }
                )
                # predicted_next_title_num = PdfUtil.predict_next_title_num(headings)

        # return headings
        headings = PdfUtil.remove_unsort_bookmarks(headings)
        return PdfUtil.complete_title(pages, headings)

    @staticmethod
    def complete_title(pages, headings: []):
        PdfUtil.log('==complete_title==')
        for i, heading in enumerate(headings):
            if not re.match(r'^[1-9]{1,2}(\.\d{1,2})*\.? ?$', heading['text']):
                continue
            heading['text'] = heading['text'].strip()
            page_layout = sorted(pages[heading['page'] - 1], key=lambda x: (-x.y0, x.x0))
            for j, element in enumerate(page_layout):
                if abs(element.y0 - heading['y']) > 20:
                    continue
                text = PdfUtil.extract_text_from_element(element).strip()
                PdfUtil.log(f'complete_title - {text}')
                if text:
                    headings[i]['text'] = headings[i]['title_num'] + ' ' + text
                    break

        return headings

    @staticmethod
    def remove_unsort_bookmarks(headings: []):
        """去除不符合排序的书签"""
        if len(headings) < 3:
            return headings

        base_pos = -1
        for i in range(len(headings) - 3):
            h0 = headings[i]
            h1 = headings[i + 1]
            h2 = headings[i + 2]
            if PdfUtil.is_same_level_and_continuous(h0, h1) and PdfUtil.is_same_level_and_continuous(h1, h2):
                base_pos = i
        PdfUtil.log(f'remove_unsort_bookmarks - base_pos={base_pos}, text=[{headings[base_pos]["text"]}]')
        if base_pos < 0:
            return headings

        new_headings = []

        lpos = base_pos - 1
        rpos = base_pos
        new_headings.insert(0, headings[rpos])
        while lpos >= 0:
            if PdfUtil.compare_title_num(headings[rpos]['title_num'], headings[lpos]['title_num']) > 0:
                new_headings.insert(0, headings[lpos])
                rpos = lpos
                lpos -= 1
            else:
                if headings[lpos]['level'] > 0 and not headings[lpos]['title_num'].endswith('1'):
                    title_num = headings[lpos]['title_num']
                    pos = title_num.rfind('.')
                    prev_title_num = f'{title_num[:pos]}.{int(title_num[pos + 1:]) - 1}'
                    dead_loop = True
                    for i in range(lpos - 1, -1, -1):
                        if headings[i]['title_num'] == prev_title_num:
                            lpos = i
                            dead_loop = False
                            break
                    if dead_loop:
                        break
                else:
                    lpos -= 1

        lpos = base_pos
        rpos = base_pos + 1
        while rpos < len(headings):
            if PdfUtil.compare_title_num(headings[rpos]['title_num'], headings[lpos]['title_num']) > 0:
                new_headings.append(headings[rpos])
                lpos = rpos
            rpos += 1

        return new_headings

    @staticmethod
    def is_same_level_and_continuous(h1, h2):
        if h1['level'] == h2['level'] and h1['level'] > 0:
            pos1 = h1['title_num'].rfind('.')
            pos2 = h2['title_num'].rfind('.')
            if h1['title_num'][:pos1] == h2['title_num'][:pos2]:
                level_index1 = int(h1['title_num'][pos1 + 1:])
                level_index2 = int(h2['title_num'][pos2 + 1:])
                if level_index1 + 1 == level_index2:
                    return True
        return False

    @staticmethod
    def compare_title_num(title_num1: str, title_num2: str):
        """比较两个标题的大小"""
        if title_num1 is None and title_num2 is None:
            return 0
        if title_num1 is None:
            return -1
        if title_num2 is None:
            return 1

        num_split1 = title_num1.split('.')
        num_split2 = title_num2.split('.')
        for i in range(min(len(num_split1), len(num_split2))):
            if num_split1[i] != num_split2[i]:
                return int(num_split1[i]) - int(num_split2[i])

        return len(num_split1) - len(num_split2)

    @staticmethod
    def remove_prev_mistake_bookmarks(headings: [], title_num: str):
        """
        删除之前误添加的标题
        例如，
        - 1.2.1 的前一个标题可能为： 1.2
        - 1.2.4 的前一个标题可能为：
          - 1.2.3
          - 1.2.3.xx
        """
        if title_num and title_num.count('.') == 0:
            return headings
        pos = title_num.rfind('.')
        level_index = int(title_num[pos + 1:])
        if level_index == 1:
            prev_title_num = title_num[:pos]
            for i in range(len(headings) - 1, -1, -1):
                if headings[i]['title_num'] == prev_title_num:
                    return headings[:i + 1]
        else:
            prev_title_num = f'{title_num[:pos]}.{int(title_num[pos + 1:]) - 1}'
            for i in range(len(headings) - 1, -1, -1):
                if headings[i]['title_num'].startswith(prev_title_num):
                    return headings[:i + 1]

        return headings

    @staticmethod
    def predict_next_title_num(headings: []):
        """
        推测下一个标题的 title。若返回 None，则表示无法推测
        举例：
            假如当前标题为 1.2.2，那么下一个标题可能为：
            - 同级：1.2.3
            - 子级：1.2.2.1
            - 父级：1.3
            - 祖父级：2
        """
        if len(headings) < 2:
            return None
        title_num1:str = headings[-2]['title_num']
        title_num2:str = headings[-1]['title_num']

        if title_num1.count('.') != title_num2.count('.'):
            # title_num2 是 title_num1 的（孙）子级标题
            if title_num2.count(title_num1) == 0:
                return None
            pos = title_num2.rfind('.')
        else:
            # title_num2 与 title_num1 是同级标题
            pos = title_num2.rfind('.')
            if title_num1[:pos] != title_num2[:pos]:
                return None
            if int(title_num1[pos + 1:]) >= int(title_num2[pos + 1:]):
                return None

        num_list = []
        prefix = title_num2[:pos]
        num_list.append(f'{prefix}.{int(title_num2[pos + 1:]) + 1}') # 同级
        num_list.append(f'{title_num2}.1') # 子级

        # 下一个父级
        while prefix:
            pos = prefix.rfind('.')
            if pos < 0:
                num_list.append(f'{int(prefix) + 1}')
                break
            else:
                num_list.append(f'{prefix[:pos]}.{int(prefix[pos + 1:]) + 1}')
                prefix = prefix[:pos]

        return num_list

    @staticmethod
    def find_parent(stack: [], level):
        parent_obj = None
        for j in range(len(stack) - 1, 0, -1):
            if stack[j]['level'] < level:
                parent_obj = stack[j]['obj']
                break
        return parent_obj

    @staticmethod
    def create_pdf_with_bookmarks(original_pdf_path, headings, output_pdf_path, progress_callback=None):
        page_total = PdfUtil.get_pdf_page_count(original_pdf_path)

        writer = PdfWriter()
        reader = PdfReader(original_pdf_path)

        # 遍历标题并添加到 PDF
        stack = [{'level': -1, 'obj': None}]
        j = 0
        for i, heading in enumerate(headings):
            while j < heading['page']:
                writer.add_page(reader.pages[j])
                j += 1
                if progress_callback:
                    progress_callback(page_total, j)

            # 添加书签
            level = heading['level']
            while stack[-1]['level'] >= level:
                stack.pop(-1)

            parent_obj = PdfUtil.find_parent(stack, level)
            obj = writer.add_outline_item(heading['text'], heading['page'] - 1, parent=parent_obj, fit=Fit.xyz(top=heading['y'] + 20))
            stack.append({'level': level, 'obj': obj})

        while j < len(reader.pages):
            writer.add_page(reader.pages[j])
            j += 1
            if progress_callback:
                progress_callback(page_total, j)

        # 保存新的 PDF 文件
        with open(output_pdf_path, "wb") as f:
            writer.write(f)